# -*- coding: utf-8 -*-
"""
Created on Mon Dec 20 20:40:24 2021

@author: wulong
"""
import numpy as np
from casadi import *
from Basic_paras import* 

# In[1] Get L-MPC solution
def get_results_l(optimalValues, pred_horzn):
    
    u = np.zeros((pred_horzn, Nuc_l))
    z = np.zeros((pred_horzn, Nuz_l))
    x = np.zeros((pred_horzn+1, Nx_l))
    y = np.zeros((pred_horzn, Ny_l))
    e = np.zeros((pred_horzn, 3))
    
    x[0,:] = optimalValues[0:Nx_l]
    
    # Create new optimal values Array
    optimal_values_new = optimalValues[Nx_l:]
    
    seed_point = Nuc_l + Nuz_l + Nx_l + Ny_l + 3
    
    i_u = 0
    i_z = Nuc_l
    i_x = Nuc_l + Nuz_l
    i_y = Nuc_l + Nuz_l + Nx_l
    i_e = Nuc_l + Nuz_l + Nx_l + Ny_l
    
    for j in range(pred_horzn):
        
        u[j,:] = optimal_values_new[int(i_u + seed_point*j):int(i_u + seed_point*j + Nuc_l)]
        z[j,:] = optimal_values_new[int(i_z + seed_point*j):int(i_z + seed_point*j + Nuz_l)]
        x[j+1,:] = optimal_values_new[int(i_x + seed_point*j):int(i_x + seed_point*j + Nx_l)]
        y[j,:] = optimal_values_new[int(i_y + seed_point*j):int(i_y + seed_point*j + Ny_l)]
        e[j,:] = optimal_values_new[int(i_e + seed_point*j):int(i_e + seed_point*j + 3)]
        
        pass
    
    return u, z, x, y, e
# In[2] Get S-MPC solution
def get_results_s(optimalValues, pred_horzn):
    
    u = np.zeros((pred_horzn, Nuc_s))
    u_m = np.zeros((pred_horzn, Nuc_m))
    u_f = np.zeros((pred_horzn, Nuc_f))
    x_m = np.zeros((pred_horzn, Nx_m))
    x_f = np.zeros((pred_horzn, Nx_f))
    x = np.zeros((pred_horzn+1, Nx_s))
    y = np.zeros((pred_horzn, Ny_s))
    yspe = np.zeros((pred_horzn, 1))
    e = np.zeros((pred_horzn, 4))
    
    x[0,:] = optimalValues[0:Nx_s]
    
    # Create new optimal values Array
    optimal_values_new = optimalValues[Nx_s:]
    
    seed_point = Nuc_s + Nuc_m + Nuc_f + Nx_m + Nx_f + Nx_s + Ny_s + 1 + 4
    
    i_u = 0
    i_um = Nuc_s
    i_uf = Nuc_s + Nuc_m
    i_xm = Nuc_s + Nuc_m + Nuc_f
    i_xf = Nuc_s + Nuc_m + Nuc_f + Nx_m
    i_x = Nuc_s + Nuc_m + Nuc_f + Nx_m + Nx_f
    i_y = Nuc_s + Nuc_m + Nuc_f + Nx_m + Nx_f + Nx_s
    i_yspe = Nuc_s + Nuc_m + Nuc_f + Nx_m + Nx_f + Nx_s + Ny_s
    i_e = Nuc_s + Nuc_m + Nuc_f + Nx_m + Nx_f + Nx_s + Ny_s + 1
    
    for j in range(pred_horzn):
        
        u[j,:] = optimal_values_new[int(i_u + seed_point*j):int(i_u + seed_point*j + Nuc_s)]
        u_m[j,:] = optimal_values_new[int(i_um + seed_point*j):int(i_um + seed_point*j + Nuc_m)]
        u_f[j,:] = optimal_values_new[int(i_uf + seed_point*j):int(i_uf + seed_point*j + Nuc_f)]
        x_m[j,:] = optimal_values_new[int(i_xm + seed_point*j):int(i_xm + seed_point*j + Nx_m)]
        x_f[j,:] = optimal_values_new[int(i_xf + seed_point*j):int(i_xf + seed_point*j + Nx_f)]
        x[j+1,:] = optimal_values_new[int(i_x + seed_point*j):int(i_x + seed_point*j + Nx_s)]
        y[j,:] = optimal_values_new[int(i_y + seed_point*j):int(i_y + seed_point*j + Ny_s)]
        yspe[j,:] = optimal_values_new[int(i_yspe + seed_point*j):int(i_yspe + seed_point*j + 1)]
        e[j,:] = optimal_values_new[int(i_e + seed_point*j):int(i_e + seed_point*j + 4)]
                
        pass
    
    return u, u_m, u_f, x_m, x_f, x, y, yspe, e

# In[3] Get M-MPC solution
def get_results_m(optimalValues, pred_horzn):
    
    u = np.zeros((pred_horzn, Nuc_m))
    u_f = np.zeros((pred_horzn, Nuc_f))
    x_f = np.zeros((pred_horzn, Nx_f))
    x = np.zeros((pred_horzn+1, Nx_m))
    y = np.zeros((pred_horzn, Ny_m))
    e = np.zeros((pred_horzn, 4))
    
    x[0,:] = optimalValues[0:Nx_m]
    
    # Create new optimal values Array
    optimal_values_new = optimalValues[Nx_m:]
    
    seed_point = Nuc_m + Nuc_f + Nx_f + Nx_m + Ny_m + 4
    
    i_u = 0
    i_uf = Nuc_m
    i_xf = Nuc_m + Nuc_f
    i_x = Nuc_m + Nuc_f + Nx_f
    i_y = Nuc_m + Nuc_f + Nx_f + Nx_m
    i_e = Nuc_m + Nuc_f + Nx_f + Nx_m + Ny_m
    
    for j in range(pred_horzn):
        
        u[j,:] = optimal_values_new[int(i_u + seed_point*j):int(i_u + seed_point*j + Nuc_m)]
        u_f[j,:] = optimal_values_new[int(i_uf + seed_point*j):int(i_uf + seed_point*j + Nuc_f)]
        x_f[j,:] = optimal_values_new[int(i_xf + seed_point*j):int(i_xf + seed_point*j + Nx_f)]
        x[j+1,:] = optimal_values_new[int(i_x + seed_point*j):int(i_x + seed_point*j + Nx_m)]
        y[j,:] = optimal_values_new[int(i_y + seed_point*j):int(i_y + seed_point*j + Ny_m)]
        e[j,:] = optimal_values_new[int(i_e + seed_point*j):int(i_e + seed_point*j + 4)]
                
        pass
    
    return u, u_f, x_f, x, y, e

# In[4] Get F-MPC solution
def get_results_f(optimalValues, pred_horzn):
    
    u = np.zeros((pred_horzn, Nuc_f))
    x = np.zeros((pred_horzn+1, Nx_f))
    y = np.zeros((pred_horzn, Ny_f))
    e = np.zeros((pred_horzn, 4))
    
    x[0,:] = optimalValues[0:Nx_f]
    
    # Create new optimal values Array
    optimal_values_new = optimalValues[Nx_f:]
    
    seed_point = Nuc_f + Nx_f + Ny_f + 4
    
    i_u = 0
    i_x = Nuc_f
    i_y = Nuc_f + Nx_f
    i_e = Nuc_f + Nx_f + Ny_f
    
    for j in range(pred_horzn):
        
        u[j,:] = optimal_values_new[int(i_u + seed_point*j):int(i_u + seed_point*j + Nuc_f)]
        x[j+1,:] = optimal_values_new[int(i_x + seed_point*j):int(i_x + seed_point*j + Nx_f)]
        y[j,:] = optimal_values_new[int(i_y + seed_point*j):int(i_y + seed_point*j + Ny_f)]
        e[j,:] = optimal_values_new[int(i_e + seed_point*j):int(i_e + seed_point*j + 4)]
        
        pass
    
    return u, x, y, e